import json
import logging
import os

from dotenv import load_dotenv
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

from catalog.dip_catalog import DipCatalog
from util.read_txt_file import ReadTxt

logger = logging.getLogger(__name__)

class MatchTcmSuperiorDiseaseGroupingBase:
    def __init__(self):
        self.tcm_knowledge = self.load_tcm_knowledge()

        # 初始化AI模型
        load_dotenv()
        api_key = os.getenv("OPENAI_API_KEY")
        base_url = os.getenv("OPENAI_BASE_URL")
        model_name = os.getenv("MODEL_NAME")

        self.llm = ChatOpenAI(
            temperature=0.1,
            max_tokens=300,
            base_url=base_url,
            model=model_name,
            api_key=api_key
        )

        # 简单的年龄分析提示
        self.tcm_prompt = ChatPromptTemplate.from_template(
            """你是中医优势病种调整专家，请对患者进行中医优势病种调整分析。

患者姓名：{patient_name}
诊断编码：{diagnosis_code}
原始分组结果：{original_result}
中医优势病种医疗知识库：{tcm_knowledge}
中医优势病种目录数据：{tcm_catalog_data}
总费用: {total_cost}
中医费用 :{tcm_cost}

请返回JSON格式的结果：

{{
  "分组名称": "调整后的治疗方式名称",
  "分值": 调整后的数值,
  "调整说明": "调整原因和计算过程",
  "中医优势病种调整完成": true
}}

**重要要求**：
1. 必须返回有效的JSON格式
2. 分值必须是数值类型，不是字符串
3. 不要添加任何JSON之外的文字

返回JSON："""
        )
        self.json_parser = JsonOutputParser()
        self.tcm_chain = self.tcm_prompt | self.llm | self.json_parser

    def ai_check_tcm_agent(self, ai_result, diagnosis_code):
        dipCatalog = DipCatalog()
        catalog_tcm = dipCatalog.load_base_catalog_excel("呼伦贝尔市DIP支付方式改革中医优势病种目录.xlsx","诊断亚目", diagnosis_code)
        try:
            tcm_res = json.loads(catalog_tcm)
        except Exception as e:
            logger.error(f"中医优势病种目录 JSON 解析失败: {e}, 原始数据: {catalog_tcm}")
            return False
        tcm_record = tcm_res.get("data", None)

        try:
            if tcm_record is None:
                return False
            # 检查解析后的结果中是否需要中医优势病种调整
            if ai_result.get('需要中医优势病种调整', False):
                return True

            return False
        except Exception as e:
            logger.error(f"AI判断中医优势病种智能体需求失败: {str(e)}")
            return False

    # 执行中医优势病种智能体
    def call_tcm_agent(self, patient_row, base_result):
        try:
            dipCatalog = DipCatalog()
            diagnosis_code = base_result.get('诊断编码', '')
            catalog_tcm = dipCatalog.load_base_catalog_excel("呼伦贝尔市DIP支付方式改革中医优势病种目录.xlsx","诊断亚目", diagnosis_code)

            result = self.tcm_chain.invoke({
                "patient_name": patient_row['人员姓名'],
                "diagnosis_code": patient_row['诊断编码'],
                "total_cost":patient_row['总费用'],
                "tcm_cost": patient_row['中医费用'],
                "original_result": base_result,
                "tcm_catalog_data": catalog_tcm,
                "tcm_knowledge": self.tcm_knowledge
            })
            logger.info(f"中医优势病种智智能体分析结果: {result}")

            # 如果返回的是字典，直接返回；如果是字符串，尝试解析JSON
            if isinstance(result, dict):
                return result
            elif isinstance(result, str):
                try:
                    return json.loads(result)
                except json.JSONDecodeError:
                    return {"错误": "JSON解析失败", "原始结果": result}
            else:
                return {"错误": "未知返回格式", "原始结果": str(result)}

        except Exception as e:
            return {"错误": str(e)}

    # 加载中医优势病种医疗知识库
    def load_tcm_knowledge(self):
        try:
            reader_txt = ReadTxt('中医优势病种医疗知识库.txt')
            return reader_txt.load_txt_lines()
        except Exception as e:
            logger.error(f"加载中医优势病种医疗知识库失败: {str(e)}")
            return ""